import tushare as ts
import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from statsmodels.tsa.stattools import coint
from sklearn.preprocessing import StandardScaler
import matplotlib.pyplot as plt
import seaborn as sns
# 设置中文字体
plt.rcParams['font.sans-serif'] = ['SimHei']  # 使用黑体显示中文
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题
# 初始化tushare pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')

def get_stock_data(ts_code, start_date, end_date):
    """获取股票日线数据"""
    df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df = df.sort_values('trade_date').set_index('trade_date')
    return df

def get_index_data(index_code, start_date, end_date):
    """获取指数日线数据"""
    df = pro.index_daily(ts_code=index_code, start_date=start_date, end_date=end_date)
    df['trade_date'] = pd.to_datetime(df['trade_date'])
    df = df.sort_values('trade_date').set_index('trade_date')
    return df

# 获取示例数据
start_date = '20200101'
end_date = '20231231'

# 获取个股数据
stock_a = get_stock_data('600519.SH', start_date, end_date)  # 贵州茅台
stock_b = get_stock_data('000858.SZ', start_date, end_date)  # 五粮液

# 获取指数数据
hs300 = get_index_data('000300.SH', start_date, end_date)    # 沪深300




#2. 量价策略 - 双均线策略
def moving_average_strategy(data, short_window=5, long_window=20):
    """双均线策略"""
    signals = pd.DataFrame(index=data.index)
    signals['price'] = data['close']
    signals['short_ma'] = data['close'].rolling(window=short_window, min_periods=1).mean()
    signals['long_ma'] = data['close'].rolling(window=long_window, min_periods=1).mean()

    # 生成交易信号
    signals['signal'] = 0.0
    signals['signal'][short_window:] = np.where(
        signals['short_ma'][short_window:] > signals['long_ma'][short_window:], 1.0, 0.0)

    # 计算实际持仓变化
    signals['positions'] = signals['signal'].diff()

    return signals


# 应用双均线策略
ma_signals = moving_average_strategy(stock_a)

# 可视化
plt.figure(figsize=(12, 6))
plt.plot(ma_signals['price'], label='Price')
plt.plot(ma_signals['short_ma'], label=f'{5}日均线')
plt.plot(ma_signals['long_ma'], label=f'{20}日均线')

# 标记买入信号
plt.plot(ma_signals.loc[ma_signals['positions'] == 1.0].index,
         ma_signals['short_ma'][ma_signals['positions'] == 1.0],
         '^', markersize=10, color='g', label='买入')

# 标记卖出信号
plt.plot(ma_signals.loc[ma_signals['positions'] == -1.0].index,
         ma_signals['short_ma'][ma_signals['positions'] == -1.0],
         'v', markersize=10, color='r', label='卖出')

plt.legend()
plt.title('双均线交易策略')
plt.show()


#3. K线形态策略 - 锤子线和吞没形态
def candlestick_patterns(data):
    """识别K线形态"""
    patterns = pd.DataFrame(index=data.index)

    # 锤子线识别
    body = abs(data['open'] - data['close'])
    upper_shadow = data['high'] - data[['open', 'close']].max(axis=1)
    lower_shadow = data[['open', 'close']].min(axis=1) - data['low']

    # 锤子线条件: 小实体，长下影线，短上影线
    hammer = (body <= 0.3 * (data['high'] - data['low'])) & \
             (lower_shadow >= 2 * body) & \
             (upper_shadow <= 0.1 * (data['high'] - data['low']))

    # 看涨吞没形态
    bullish_engulfing = (data['close'] > data['open']) & \
                        (data['close'].shift(1) < data['open'].shift(1)) & \
                        (data['close'] > data['open'].shift(1)) & \
                        (data['open'] < data['close'].shift(1))

    patterns['hammer'] = hammer
    patterns['bullish_engulfing'] = bullish_engulfing

    return patterns


# 应用K线形态识别
patterns = candlestick_patterns(stock_a)

# 可视化示例
sample_data = stock_a.iloc[50:70].copy()
sample_patterns = candlestick_patterns(sample_data)

plt.figure(figsize=(12, 6))
for i in range(len(sample_data)):
    color = 'r' if sample_data['close'].iloc[i] < sample_data['open'].iloc[i] else 'g'
    plt.plot([i, i], [sample_data['low'].iloc[i], sample_data['high'].iloc[i]], color=color)
    plt.plot([i, i], [sample_data['open'].iloc[i], sample_data['close'].iloc[i]],
             linewidth=5, color=color)

    # 标记锤子线
    if sample_patterns['hammer'].iloc[i]:
        plt.text(i, sample_data['high'].iloc[i], 'H', ha='center', va='bottom', fontsize=12)

    # 标记吞没形态
    if sample_patterns['bullish_engulfing'].iloc[i]:
        plt.text(i, sample_data['high'].iloc[i], 'BE', ha='center', va='bottom', fontsize=12)

plt.xticks(range(len(sample_data)), sample_data.index.date, rotation=45)
plt.title('K线形态识别示例')
plt.show()

#4. 统计套利策略 - 配对交易
def find_cointegrated_pairs(price_data):
    """寻找协整的股票对"""
    n = price_data.shape[1]
    score_matrix = np.zeros((n, n))
    pvalue_matrix = np.ones((n, n))
    pairs = []

    for i in range(n):
        for j in range(i + 1, n):
            S1 = price_data.iloc[:, i]
            S2 = price_data.iloc[:, j]
            _, pvalue, _ = coint(S1, S2)
            pvalue_matrix[i, j] = pvalue
            if pvalue < 0.05:  # 显著性水平5%
                pairs.append((price_data.columns[i], price_data.columns[j], pvalue))

    return pvalue_matrix, pairs


def pairs_trading_strategy(stock1, stock2, lookback=20):
    """配对交易策略"""
    df = pd.DataFrame()
    df['stock1'] = stock1['close']
    df['stock2'] = stock2['close']

    # 计算价差
    df['spread'] = df['stock1'] - df['stock2']

    # 计算z-score
    df['mean'] = df['spread'].rolling(window=lookback).mean()
    df['std'] = df['spread'].rolling(window=lookback).std()
    df['zscore'] = (df['spread'] - df['mean']) / df['std']

    # 生成交易信号
    df['long_entry'] = df['zscore'] <= -2.0
    df['long_exit'] = df['zscore'] >= 0.0
    df['short_entry'] = df['zscore'] >= 2.0
    df['short_exit'] = df['zscore'] <= 0.0

    # 初始化仓位
    df['positions'] = 0

    # 生成交易信号
    long_position = False
    short_position = False

    for i in range(len(df)):
        if df['long_entry'].iloc[i] and not long_position:
            df.loc[df.index[i], 'positions'] = 1  # 买入stock1，卖出stock2
            long_position = True
        elif df['long_exit'].iloc[i] and long_position:
            df.loc[df.index[i], 'positions'] = 0
            long_position = False
        elif df['short_entry'].iloc[i] and not short_position:
            df.loc[df.index[i], 'positions'] = -1  # 卖出stock1，买入stock2
            short_position = True
        elif df['short_exit'].iloc[i] and short_position:
            df.loc[df.index[i], 'positions'] = 0
            short_position = False

    return df




# 寻找协整股票对
price_data = pd.concat([stock_a['close'], stock_b['close']], axis=1)
price_data.columns = ['stock_a', 'stock_b']
_, pairs = find_cointegrated_pairs(price_data)
from statsmodels.tsa.stattools import adfuller

def check_stationarity(data):
    """检查时间序列的平稳性"""
    result = adfuller(data)
    print(f"ADF Statistic: {result[0]}")
    print(f"p-value: {result[1]}")
    return result[1] < 0.05
stock_a_diff = stock_a['close'].diff().dropna()
stock_b_diff = stock_b['close'].diff().dropna()
price_data_diff = pd.concat([stock_a_diff, stock_b_diff], axis=1)
price_data_diff.columns = ['stock_a', 'stock_b']
_, pairs = find_cointegrated_pairs(price_data_diff)
print("找到的协整股票对:", pairs)

# 应用配对交易策略
if pairs:
    pt_signals = pairs_trading_strategy(stock_a, stock_b)

    # 可视化
    plt.figure(figsize=(14, 7))

    # 价差和z-score
    plt.subplot(211)
    plt.plot(pt_signals['spread'], label='价差')
    plt.axhline(pt_signals['mean'].iloc[-1], color='gray', linestyle='--')
    plt.axhline(pt_signals['mean'].iloc[-1] + 2 * pt_signals['std'].iloc[-1], color='r', linestyle='--')
    plt.axhline(pt_signals['mean'].iloc[-1] - 2 * pt_signals['std'].iloc[-1], color='g', linestyle='--')
    plt.legend()
    plt.title('价差和标准差通道')

    # 交易信号
    plt.subplot(212)
    plt.plot(pt_signals['zscore'], label='Z-score')
    plt.axhline(2, color='r', linestyle='--')
    plt.axhline(-2, color='g', linestyle='--')
    plt.axhline(0, color='gray', linestyle='--')

    # 标记交易信号
    plt.plot(pt_signals[pt_signals['positions'] == 1].index,
             pt_signals['zscore'][pt_signals['positions'] == 1],
             '^', markersize=10, color='g', label='买入信号')
    plt.plot(pt_signals[pt_signals['positions'] == -1].index,
             pt_signals['zscore'][pt_signals['positions'] == -1],
             'v', markersize=10, color='r', label='卖出信号')

    plt.legend()
    plt.title('Z-score和交易信号')
    plt.tight_layout()
    plt.show()

import pandas as pd
import numpy as np


def add_technical_indicators(data):
    """添加技术指标"""
    # 移动平均线
    data['MA5'] = data['close'].rolling(window=5).mean()
    data['MA20'] = data['close'].rolling(window=20).mean()

    # 相对强弱指数 (RSI)
    delta = data['close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    data['RSI'] = 100 - (100 / (1 + gain / loss))

    # 布林带
    data['BB_upper'] = data['close'].rolling(window=20).mean() + 2 * data['close'].rolling(window=20).std()
    data['BB_lower'] = data['close'].rolling(window=20).mean() - 2 * data['close'].rolling(window=20).std()

    # 去掉计算指标时产生的NaN值
    data = data.dropna()
    return data


# 添加技术指标
stock_a = add_technical_indicators(stock_a)
stock_b = add_technical_indicators(stock_b)

from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error


def train_predict_model(data, target_col='close', features=None, days=5):
    """训练和预测模型"""
    if features is None:
        features = ['MA5', 'MA20', 'RSI', 'BB_upper', 'BB_lower']

    # 构建特征和目标
    X = []
    y = []
    for i in range(days, len(data)):
        X.append(data[features].iloc[i - days:i].values.flatten())
        y.append(data[target_col].iloc[i])

    X = np.array(X)
    y = np.array(y)

    # 划分训练集和测试集
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

    # 训练模型
    model = LinearRegression()
    model.fit(X_train, y_train)

    # 预测
    y_pred = model.predict(X_test)

    # 评估模型
    mse = mean_squared_error(y_test, y_pred)
    print(f"均方误差 (MSE): {mse}")

    return model, X_test, y_test, y_pred


# 训练模型
model_a, X_test_a, y_test_a, y_pred_a = train_predict_model(stock_a)
model_b, X_test_b, y_test_b, y_pred_b = train_predict_model(stock_b)

# 可视化预测结果
plt.figure(figsize=(12, 6))
plt.plot(y_test_a, label='实际价格')
plt.plot(y_pred_a, label='预测价格', linestyle='--')
plt.legend()
plt.title(' 贵州茅台价格预测')
plt.show()

plt.figure(figsize=(12, 6))
plt.plot(y_test_b, label='实际价格')
plt.plot(y_pred_b, label='预测价格', linestyle='--')
plt.legend()
plt.title(' 五粮液价格预测')
plt.show()


def ml_trading_strategy(data, model, features=None, days=5):
    """基于机器学习的交易策略"""
    if features is None:
        features = ['MA5', 'MA20', 'RSI', 'BB_upper', 'BB_lower']

    signals = pd.DataFrame(index=data.index)
    signals['price'] = data['close']
    signals['signal'] = 0.0

    for i in range(days, len(data)):
        X = data[features].iloc[i - days:i].values.flatten().reshape(1, -1)
        predicted_price = model.predict(X)
        if predicted_price > data['close'].iloc[i - 1]:
            signals.loc[data.index[i], 'signal'] = 1.0  # 买入
        elif predicted_price < data['close'].iloc[i - 1]:
            signals.loc[data.index[i], 'signal'] = -1.0  # 卖出

    signals['positions'] = signals['signal'].diff()
    return signals


# 应用交易策略
ml_signals_a = ml_trading_strategy(stock_a, model_a)
ml_signals_b = ml_trading_strategy(stock_b, model_b)

# 可视化交易信号
plt.figure(figsize=(12, 6))
plt.plot(ml_signals_a['price'], label='Price')
plt.plot(ml_signals_a['price'][ml_signals_a['positions'] == 1], '^', markersize=10, color='g', label='买入')
plt.plot(ml_signals_a['price'][ml_signals_a['positions'] == -1], 'v', markersize=10, color='r', label='卖出')
plt.legend()
plt.title('基于机器学习的交易策略 - 贵州茅台')
plt.show()

plt.figure(figsize=(12, 6))
plt.plot(ml_signals_b['price'], label='Price')
plt.plot(ml_signals_b['price'][ml_signals_b['positions'] == 1], '^', markersize=10, color='g', label='买入')
plt.plot(ml_signals_b['price'][ml_signals_b['positions'] == -1], 'v', markersize=10, color='r', label='卖出')
plt.legend()
plt.title('基于机器学习的交易策略 - 五粮液')
plt.show()



#2
class QuantTradingSystem:
    def __init__(self, capital=1000000, risk_per_trade=0.02, stop_loss=0.05, take_profit=0.10):
        """
        初始化量化交易系统
        :param capital: 初始资金
        :param risk_per_trade: 每笔交易风险比例
        :param stop_loss: 止损比例
        :param take_profit: 止盈比例
        """
        self.initial_capital = capital
        self.capital = capital
        self.risk_per_trade = risk_per_trade
        self.stop_loss = stop_loss
        self.take_profit = take_profit
        self.positions = {}
        self.trade_history = []
        self.stock_data = {}

    def get_stock_data(self, ts_codes, start_date, end_date):
        """获取多只股票数据"""
        for ts_code in ts_codes:
            df = pro.daily(ts_code=ts_code, start_date=start_date, end_date=end_date)
            df['trade_date'] = pd.to_datetime(df['trade_date'])
            df = df.sort_values('trade_date').set_index('trade_date')
            self.stock_data[ts_code] = self.add_technical_indicators(df)

    def add_technical_indicators(self, data):
        """添加技术指标"""
        # 移动平均线
        data['MA5'] = data['close'].rolling(window=5).mean()
        data['MA20'] = data['close'].rolling(window=20).mean()

        # 相对强弱指数 (RSI)
        delta = data['close'].diff()
        gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
        loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
        data['RSI'] = 100 - (100 / (1 + gain / loss))

        # 布林带
        data['BB_upper'] = data['close'].rolling(window=20).mean() + 2 * data['close'].rolling(window=20).std()
        data['BB_lower'] = data['close'].rolling(window=20).mean() - 2 * data['close'].rolling(window=20).std()

        # MACD
        data['EMA12'] = data['close'].ewm(span=12, adjust=False).mean()
        data['EMA26'] = data['close'].ewm(span=26, adjust=False).mean()
        data['MACD'] = data['EMA12'] - data['EMA26']
        data['Signal'] = data['MACD'].ewm(span=9, adjust=False).mean()

        # 去掉计算指标时产生的NaN值
        data = data.dropna()
        return data

    def generate_signals(self, ts_code):
        """生成交易信号"""
        data = self.stock_data[ts_code].copy()

        # 1. 双均线策略信号
        data['MA_signal'] = np.where(data['MA5'] > data['MA20'], 1, 0)

        # 2. RSI策略信号
        data['RSI_signal'] = 0
        data.loc[data['RSI'] < 30, 'RSI_signal'] = 1  # 超卖，买入信号
        data.loc[data['RSI'] > 70, 'RSI_signal'] = -1  # 超买，卖出信号

        # 3. 布林带策略信号
        data['BB_signal'] = 0
        data.loc[data['close'] < data['BB_lower'], 'BB_signal'] = 1  # 低于下轨，买入
        data.loc[data['close'] > data['BB_upper'], 'BB_signal'] = -1  # 高于上轨，卖出

        # 4. MACD策略信号
        data['MACD_signal'] = 0
        data.loc[
            (data['MACD'] > data['Signal']) & (data['MACD'].shift(1) <= data['Signal'].shift(1)), 'MACD_signal'] = 1
        data.loc[
            (data['MACD'] < data['Signal']) & (data['MACD'].shift(1) >= data['Signal'].shift(1)), 'MACD_signal'] = -1

        # 综合信号
        data['signal'] = data['MA_signal'] + data['RSI_signal'] + data['BB_signal'] + data['MACD_signal']

        # 信号平滑处理
        data['signal'] = np.where(data['signal'] >= 2, 1, np.where(data['signal'] <= -2, -1, 0))

        # 标记实际交易点
        data['positions'] = data['signal'].diff()

        self.stock_data[ts_code] = data
        return data

    def execute_trades(self, ts_code, date):
        """执行交易"""
        if ts_code not in self.stock_data:
            return

        data = self.stock_data[ts_code]
        if date not in data.index:
            return

        row = data.loc[date]
        price = row['close']

        # 检查是否有持仓
        position = self.positions.get(ts_code, {'shares': 0, 'entry_price': 0})

        # 止损/止盈检查
        if position['shares'] > 0:
            # 止损检查
            if price <= position['entry_price'] * (1 - self.stop_loss):
                self.close_position(ts_code, date, price, '止损')
                return
            # 止盈检查
            if price >= position['entry_price'] * (1 + self.take_profit):
                self.close_position(ts_code, date, price, '止盈')
                return

        # 生成新信号
        if row['positions'] == 1 and position['shares'] == 0:  # 买入信号且无持仓
            self.open_position(ts_code, date, price)
        elif row['positions'] == -1 and position['shares'] > 0:  # 卖出信号且有持仓
            self.close_position(ts_code, date, price, '信号卖出')

    def open_position(self, ts_code, date, price):
        """开仓"""
        # 计算可购买股数
        risk_amount = self.capital * self.risk_per_trade
        max_shares = int(risk_amount / (price * self.stop_loss))

        if max_shares == 0:
            return

        # 计算实际购买股数（假设可以部分购买）
        cost = max_shares * price
        if cost > self.capital:
            max_shares = int(self.capital / price)
            cost = max_shares * price

        if max_shares == 0:
            return

        self.capital -= cost
        self.positions[ts_code] = {'shares': max_shares, 'entry_price': price, 'entry_date': date}

        # 记录交易
        self.trade_history.append({
            'date': date,
            'ts_code': ts_code,
            'action': '买入',
            'shares': max_shares,
            'price': price,
            'value': cost,
            'capital': self.capital,
            'reason': '信号买入'
        })

    def close_position(self, ts_code, date, price, reason):
        """平仓"""
        position = self.positions.get(ts_code)
        if not position or position['shares'] == 0:
            return

        # 计算平仓价值
        value = position['shares'] * price
        self.capital += value

        # 计算盈亏
        pnl = value - (position['shares'] * position['entry_price'])
        pnl_pct = pnl / (position['shares'] * position['entry_price'])

        # 记录交易
        self.trade_history.append({
            'date': date,
            'ts_code': ts_code,
            'action': '卖出',
            'shares': position['shares'],
            'price': price,
            'value': value,
            'capital': self.capital,
            'pnl': pnl,
            'pnl_pct': pnl_pct,
            'reason': reason
        })

        # 清除持仓
        self.positions[ts_code] = {'shares': 0, 'entry_price': 0}

    def backtest(self, ts_codes, start_date, end_date):
        """回测"""
        # 获取数据
        self.get_stock_data(ts_codes, start_date, end_date)

        # 生成信号
        for ts_code in ts_codes:
            self.generate_signals(ts_code)

        # 获取所有交易日
        all_dates = set()
        for ts_code in ts_codes:
            all_dates.update(self.stock_data[ts_code].index)
        all_dates = sorted(all_dates)

        # 按日期执行回测
        for date in all_dates:
            for ts_code in ts_codes:
                self.execute_trades(ts_code, date)

        # 计算回测结果
        self.calculate_performance()

        # 输出结果
        self.report_results()

    def calculate_performance(self):
        """计算回测绩效"""
        if not self.trade_history:
            return

        # 创建交易记录DataFrame
        trades = pd.DataFrame(self.trade_history)

        # 计算累计收益率
        initial_value = self.initial_capital
        final_value = self.capital

        # 加上未平仓头寸的价值
        for ts_code, pos in self.positions.items():
            if pos['shares'] > 0:
                last_price = self.stock_data[ts_code].iloc[-1]['close']
                final_value += pos['shares'] * last_price

        self.total_return = (final_value - initial_value) / initial_value
        self.annualized_return = (1 + self.total_return) ** (252 / len(trades)) - 1

        # 计算胜率
        if len(trades) > 0:
            winning_trades = trades[trades['pnl'] > 0]
            self.win_rate = len(winning_trades) / len(trades)
        else:
            self.win_rate = 0

        # 计算最大回撤
        equity_curve = [self.initial_capital]
        for trade in self.trade_history:
            equity_curve.append(trade['capital'])

        max_drawdown = 0
        peak = equity_curve[0]
        for value in equity_curve:
            if value > peak:
                peak = value
            dd = (peak - value) / peak
            if dd > max_drawdown:
                max_drawdown = dd

        self.max_drawdown = max_drawdown

    def report_results(self):
        """输出回测结果"""
        print("\n========== 回测结果 ==========")
        print(f"初始资金: {self.initial_capital:,.2f}元")
        print(f"最终资金: {self.capital:,.2f}元")
        print(f"总收益率: {self.total_return * 100:.2f}%")
        print(f"年化收益率: {self.annualized_return * 100:.2f}%")
        print(f"胜率: {self.win_rate * 100:.2f}%")
        print(f"最大回撤: {self.max_drawdown * 100:.2f}%")
        print("=============================")

        # 绘制资金曲线
        self.plot_equity_curve()

        # 显示交易记录
        self.show_trade_history()

    def plot_equity_curve(self):
        """绘制资金曲线"""
        equity = [self.initial_capital]
        dates = []

        for trade in self.trade_history:
            equity.append(trade['capital'])
            dates.append(trade['date'])

        plt.figure(figsize=(12, 6))
        plt.plot(dates, equity[1:], label='资金曲线')
        plt.title('资金曲线')
        plt.xlabel('日期')
        plt.ylabel('资金 (元)')
        plt.grid(True)
        plt.legend()
        plt.show()

    def show_trade_history(self):
        """显示交易记录"""
        trades = pd.DataFrame(self.trade_history)
        if len(trades) > 0:
            print("\n最近5笔交易记录:")
            print(trades[['date', 'ts_code', 'action', 'price', 'shares', 'pnl_pct', 'reason']].tail(5))

            # 按股票统计
            print("\n按股票统计:")
            stats = trades.groupby('ts_code').agg({
                'pnl': 'sum',
                'pnl_pct': 'mean',
                'action': 'count'
            }).rename(columns={'action': '交易次数', 'pnl': '总盈亏', 'pnl_pct': '平均收益率'})
            print(stats)


# 测试量化交易系统
if __name__ == "__main__":
    # 选择一组相关性较高的股票
    stock_list = ['600519.SH', '000858.SZ', '600809.SH', '000568.SZ']  # 贵州茅台、五粮液、山西汾酒、泸州老窖

    # 设置回测日期
    start_date = '20200101'
    end_date = '20231231'

    # 创建并运行量化交易系统
    qts = QuantTradingSystem(capital=1000000, risk_per_trade=0.02, stop_loss=0.05, take_profit=0.10)
    qts.backtest(stock_list, start_date, end_date)
